#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "lexer.h"

typedef struct ASTNode {
    char kind[32];
    char text[64];
    struct ASTNode **children;
    int child_count;
} ASTNode;

ASTNode *new_node(const char *kind, const char *text) {
    ASTNode *node = (ASTNode*)malloc(sizeof(ASTNode));
    if (!node) {
        fprintf(stderr, "Out of memory\n");
        exit(1);
    }
    memset(node->kind, 0, sizeof(node->kind));
    memset(node->text, 0, sizeof(node->text));
    if (kind) {
        strncpy(node->kind, kind, sizeof(node->kind) - 1);
    }
    if (text) {
        strncpy(node->text, text, sizeof(node->text) - 1);
    }
    node->children = NULL;
    node->child_count = 0;
    return node;
}

void add_child(ASTNode *parent, ASTNode *child) {
    if (!parent || !child) return;
    ASTNode **new_children = (ASTNode**)realloc(
        parent->children,
        sizeof(ASTNode*) * (parent->child_count + 1)
    );
    if (!new_children) {
        fprintf(stderr, "Out of memory\n");
        exit(1);
    }
    parent->children = new_children;
    parent->children[parent->child_count] = child;
    parent->child_count += 1;
}

void print_ast(ASTNode *node, int indent) {
    if (!node) return;
    for (int i = 0; i < indent; i++) {
        printf("  ");
    }
    printf("%s", node->kind);
    if (node->text[0]) {
        printf(" (%s)", node->text);
    }
    printf("\n");
    for (int i = 0; i < node->child_count; i++) {
        print_ast(node->children[i], indent + 1);
    }
}

Token cur_tok;
Token peek_tok;

void next_token() {
    cur_tok = peek_tok;
    peek_tok = get_next_token();
}

int is_type_token(const char *t) {
    return strcmp(cur_tok.type, t) == 0;
}

void syntax_error(const char *msg) {
    fprintf(stderr,
        "Syntax error at line %d, column %d: %s (token: type=%s, value=%s)\n",
        cur_tok.line, cur_tok.column, msg, cur_tok.type, cur_tok.value);
    exit(1);
}

void match_token(const char *t) {
    if (is_type_token(t)) {
        next_token();
    } else {
        char buf[128];
        snprintf(buf, sizeof(buf), "expect '%s'", t);
        syntax_error(buf);
    }
}

void match_identifier_token() {
    if (is_type_token("IDENTIFIER")) {
        next_token();
    } else {
        syntax_error("expect IDENTIFIER");
    }
}

void match_number_token() {
    if (is_type_token("NUMBER")) {
        next_token();
    } else {
        syntax_error("expect NUMBER");
    }
}

ASTNode *program();
ASTNode *fun_declaration();
ASTNode *main_declaration();
ASTNode *fun_body();
void declaration_list(ASTNode *block);
ASTNode *declaration_stat();
void statement_list(ASTNode *block);
ASTNode *statement();
ASTNode *if_stat();
ASTNode *while_stat();
ASTNode *for_stat();
ASTNode *write_stat();
ASTNode *read_stat();
ASTNode *compound_stat();
ASTNode *expression_stat();

ASTNode *expression();
ASTNode *bool_expr();
ASTNode *additive_expr();
ASTNode *term();
ASTNode *factor();

int is_statement_start() {
    return is_type_token("if") ||
           is_type_token("while") ||
           is_type_token("for") ||
           is_type_token("read") ||
           is_type_token("write") ||
           is_type_token("{") ||
           is_type_token("IDENTIFIER") ||
           is_type_token("NUMBER") ||
           is_type_token("(");
}

ASTNode *type_node() {
    if (is_type_token("int") || is_type_token("IDENTIFIER")) {
        ASTNode *t = new_node("Type", cur_tok.value);
        next_token();
        return t;
    }
    syntax_error("expect type");
    return NULL;
}

ASTNode *param() {
    if (!is_type_token("IDENTIFIER")) {
        syntax_error("expect parameter name");
    }
    ASTNode *id = new_node("Identifier", cur_tok.value);
    match_identifier_token();
    match_token(":");
    ASTNode *t = type_node();
    ASTNode *p = new_node("Param", "");
    add_child(p, id);
    add_child(p, t);
    return p;
}

void param_list(ASTNode *func_node) {
    if (!is_type_token("IDENTIFIER")) {
        return;
    }
    ASTNode *p = param();
    add_child(func_node, p);
    while (is_type_token(",")) {
        match_token(",");
        p = param();
        add_child(func_node, p);
    }
}

ASTNode *fun_declaration() {
    if (!is_type_token("func")) {
        return NULL;
    }
    match_token("func");
    if (!is_type_token("IDENTIFIER")) {
        syntax_error("expect function name after 'func'");
    }
    ASTNode *func = new_node("FuncDecl", cur_tok.value);
    match_identifier_token();
    match_token("(");
    param_list(func);
    match_token(")");
    if (is_type_token(":")) {
        match_token(":");
        ASTNode *ret = type_node();
        add_child(func, ret);
    }
    ASTNode *body = fun_body();
    add_child(func, body);
    return func;
}

ASTNode *main_declaration() {
    if (!is_type_token("main")) {
        syntax_error("program must have main");
    }
    match_token("main");
    ASTNode *main_node = new_node("MainFunc", "main");
    match_token("(");
    param_list(main_node);
    match_token(")");
    if (is_type_token(":")) {
        match_token(":");
        ASTNode *ret = type_node();
        add_child(main_node, ret);
    }
    ASTNode *body = fun_body();
    add_child(main_node, body);
    return main_node;
}

ASTNode *fun_body() {
    match_token("{");
    ASTNode *block = new_node("Block", "");
    declaration_list(block);
    statement_list(block);
    match_token("}");
    return block;
}

void declaration_list(ASTNode *block) {
    while (is_type_token("var")) {
        ASTNode *d = declaration_stat();
        add_child(block, d);
    }
}

ASTNode *declaration_stat() {
    match_token("var");
    if (!is_type_token("IDENTIFIER")) {
        syntax_error("expect identifier after 'var'");
    }
    char name[MAX_TOKEN_LEN];
    strcpy(name, cur_tok.value);
    match_identifier_token();
    match_token(":");
    if (!is_type_token("int")) {
        syntax_error("only 'int' type is supported in declaration_stat");
    }
    match_token("int");
    ASTNode *decl = new_node("VarDecl", name);
    ASTNode *type = new_node("Type", "int");
    add_child(decl, type);
    if (is_type_token("=")) {
        match_token("=");
        ASTNode *init = expression();
        add_child(decl, init);
    }
    return decl;
}


void statement_list(ASTNode *block) {
    while (is_statement_start()) {
        ASTNode *s = statement();
        add_child(block, s);
    }
}

ASTNode *statement() {
    if (is_type_token("if")) {
        return if_stat();
    } else if (is_type_token("while")) {
        return while_stat();
    } else if (is_type_token("for")) {
        return for_stat();
    } else if (is_type_token("read")) {
        return read_stat();
    } else if (is_type_token("write")) {
        return write_stat();
    } else if (is_type_token("{")) {
        return compound_stat();
    } else {
        return expression_stat();
    }
}

ASTNode *if_stat() {
    match_token("if");
    match_token("(");
    ASTNode *cond = expression();
    match_token(")");
    ASTNode *then_stmt = statement();
    ASTNode *node = new_node("If", "");
    add_child(node, cond);
    add_child(node, then_stmt);
    if (is_type_token("else")) {
        match_token("else");
        ASTNode *else_stmt = statement();
        add_child(node, else_stmt);
    }
    return node;
}

ASTNode *while_stat() {
    match_token("while");
    match_token("(");
    ASTNode *cond = expression();
    match_token(")");
    ASTNode *body = statement();
    ASTNode *node = new_node("While", "");
    add_child(node, cond);
    add_child(node, body);
    return node;
}

ASTNode *for_stat() {
    match_token("for");
    match_token("(");
    if (!is_type_token("IDENTIFIER")) {
        syntax_error("expect identifier after 'for('");
    }
    ASTNode *id = new_node("Identifier", cur_tok.value);
    match_identifier_token();
    if (!is_type_token("in")) {
        syntax_error("expect 'in' in for statement");
    }
    match_token("in");
    ASTNode *iter = expression();
    match_token(")");
    ASTNode *body = statement();
    ASTNode *node = new_node("For", "");
    add_child(node, id);
    add_child(node, iter);
    add_child(node, body);
    return node;
}

ASTNode *write_stat() {
    match_token("write");
    ASTNode *expr = expression();
    ASTNode *node = new_node("Write", "");
    add_child(node, expr);
    return node;
}

ASTNode *read_stat() {
    match_token("read");
    if (!is_type_token("IDENTIFIER")) {
        syntax_error("expect identifier after 'read'");
    }
    ASTNode *id = new_node("Identifier", cur_tok.value);
    match_identifier_token();
    ASTNode *node = new_node("Read", "");
    add_child(node, id);
    return node;
}

ASTNode *compound_stat() {
    match_token("{");
    ASTNode *block = new_node("Block", "");
    statement_list(block);
    match_token("}");
    return block;
}

ASTNode *expression_stat() {
    ASTNode *expr = expression();
    return expr;
}

ASTNode *expression() {
    if (strcmp(cur_tok.type, "IDENTIFIER") == 0 &&
        strcmp(peek_tok.type, "=") == 0) {
        ASTNode *left = new_node("Identifier", cur_tok.value);
        match_identifier_token();
        match_token("=");
        ASTNode *right = bool_expr();
        ASTNode *node = new_node("Assign", "=");
        add_child(node, left);
        add_child(node, right);
        return node;
    } else {
        return bool_expr();
    }
}

ASTNode *bool_expr() {
    ASTNode *left = additive_expr();
    if (is_type_token(">") || is_type_token("<") ||
        is_type_token(">=") || is_type_token("<=") ||
        is_type_token("==") || is_type_token("!=")) {
        char op[4];
        memset(op, 0, sizeof(op));
        strncpy(op, cur_tok.type, sizeof(op) - 1);
        next_token();
        ASTNode *right = additive_expr();
        ASTNode *node = new_node("Binary", op);
        add_child(node, left);
        add_child(node, right);
        return node;
    }
    return left;
}

ASTNode *additive_expr() {
    ASTNode *left = term();
    while (is_type_token("+") || is_type_token("-")) {
        char op[2];
        memset(op, 0, sizeof(op));
        strncpy(op, cur_tok.type, sizeof(op) - 1);
        next_token();
        ASTNode *right = term();
        ASTNode *node = new_node("Binary", op);
        add_child(node, left);
        add_child(node, right);
        left = node;
    }
    return left;
}

ASTNode *term() {
    ASTNode *left = factor();
    while (is_type_token("*") || is_type_token("/")) {
        char op[2];
        memset(op, 0, sizeof(op));
        strncpy(op, cur_tok.type, sizeof(op) - 1);
        next_token();
        ASTNode *right = factor();
        ASTNode *node = new_node("Binary", op);
        add_child(node, left);
        add_child(node, right);
        left = node;
    }
    return left;
}

ASTNode *factor() {
    if (is_type_token("(")) {
        match_token("(");
        ASTNode *node = additive_expr();
        match_token(")");
        return node;
    } else if (is_type_token("IDENTIFIER")) {
        ASTNode *id = new_node("Identifier", cur_tok.value);
        match_identifier_token();
        return id;
    } else if (is_type_token("NUMBER")) {
        ASTNode *num = new_node("Number", cur_tok.value);
        match_number_token();
        return num;
    } else {
        syntax_error("expect '(', IDENTIFIER or NUMBER in factor");
        return NULL;
    }
}

ASTNode *program() {
    ASTNode *prog = new_node("Program", "");
    while (is_type_token("func")) {
        ASTNode *f = fun_declaration();
        add_child(prog, f);
    }
    ASTNode *main_node = main_declaration();
    add_child(prog, main_node);
    return prog;
}

int main(int argc, char *argv[]) {
    if (argc < 2) {
        printf("Usage: %s <source_file>\n", argv[0]);
        return 1;
    }

    init_lexer(argv[1]);
    cur_tok = get_next_token();
    peek_tok = get_next_token();

    ASTNode *root = program();
    printf("Syntax tree:\n");
    print_ast(root, 0);

    return 0;
}
